#include "StdAfx.h"
#include "inttypes.h"
#include "FFCore.h"

#include "MediaFrame.h"
#include "MediaBuffer.h"
#include <strsafe.h>


CFFCore::CFFCore()	
	:	m_pAudioOutBuf(NULL)
	,	m_pVidoeOutBuf(NULL) 
	,	m_pAudioFifoBuf(NULL)
	,   m_pSetting(NULL)
	,	m_pFileOutFormat(NULL)
	,	m_pFormatContext(NULL)
	,	m_pFormatParameter(NULL)
	,	m_pAudioSt(NULL)
	,	m_pVideoSt(NULL)	
	,	m_pProp(NULL)
	,	fifo(NULL)
	,	m_bInit(false)
	,	m_pMediaFile(NULL)
	,	m_pImgRmp(NULL)
	,	m_pAudRmp(NULL)	
	,	m_fCreateVideoStream(TRUE)
	,	m_fCreateAudioStream(TRUE)
{	
}

CFFCore::~CFFCore(void)
{
	if(m_pVidoeOutBuf) {
		av_free(m_pVidoeOutBuf);
		m_pVidoeOutBuf = NULL;
	}
	if(m_pAudioOutBuf){
		av_free(m_pAudioOutBuf);
		m_pAudioOutBuf = NULL;
	}
	if(m_pAudioFifoBuf)
	{
		av_free(m_pAudioFifoBuf);
		m_pAudioFifoBuf = NULL;
	}
	if(fifo){
		av_fifo_free(fifo);
		av_free(fifo);	
	}
	if(m_pMediaFile)
	{
		delete m_pMediaFile;
	}
	if(m_pImgRmp) delete m_pImgRmp;
	if(m_pAudRmp) delete m_pAudRmp;
}

bool CFFCore::SetOption(COptionParameter* pOption)
{
	if(!pOption) throw "Null pointer!";
	m_pSetting = pOption;
	
	return true;
}

bool CFFCore::SetSink(char* file)
{
	if(!file) return false;
	if(m_pMediaFile) delete m_pMediaFile;
	m_pMediaFile = new CPathA;	
	m_pMediaFile->Append(file);
	return true;

}

bool CFFCore::SetSrcProp(CMediaProp* pProp)
{
	if(!pProp) return false;
	m_pProp = pProp;
	return true;
}


void CFFCore::CreateVideoStream(BOOL bVideo)
{
	m_fCreateVideoStream = bVideo;
}

void CFFCore::CreateAudioStream(BOOL bAudio)
{
	m_fCreateAudioStream = bAudio;
}

void CFFCore::checkExt()
{
	if(!m_pMediaFile) return;
	if(!m_pSetting)return;	
	if(m_pMediaFile->GetExtension().Compare(m_pSetting->ext) == 0 ) {
				
	}else{
		m_pMediaFile->RemoveExtension();
		m_pMediaFile->AddExtension(m_pSetting->ext);
	}
}

bool CFFCore::InitFfmpeg()
{
	if(!m_pMediaFile  || !m_pSetting) throw "Null pointer!";

	checkExt();
	m_nFTP = int64_t(float(m_pSetting->m_nFrameRateBase)/float(m_pSetting->m_nFrameRate) * 10000000.0 + 0.5);	
	m_nTotalFrameNumber = 0;	
	m_pFileOutFormat = guess_format(m_pSetting->format, NULL, NULL);
	if(!m_pFileOutFormat) throw "The output format is not support!";	
	fifo = (AVFifoBuffer*)av_malloc(sizeof(AVFifoBuffer));
	if(!fifo) throw "Memory Allocate failed! ";
	memset(fifo, 0, sizeof(AVFifoBuffer));
	av_fifo_init(fifo, MAX_AUDIO_OUT_SIZE);
	m_pFormatContext = av_alloc_format_context();
	if(!m_pFormatContext) throw "There is some error about format context!";
	_snprintf(m_pFormatContext->filename, sizeof(m_pFormatContext->filename), "%s", m_pMediaFile->m_strPath);
	m_pFormatContext->oformat = m_pFileOutFormat;
	CodecID vid, aid;
	m_pSetting->m_videoCodecID == CODEC_ID_NONE ? vid = m_pFileOutFormat->video_codec : vid = m_pSetting->m_videoCodecID;
	m_pSetting->m_audioCodecID == CODEC_ID_NONE ? aid = m_pFileOutFormat->audio_codec : aid = m_pSetting->m_audioCodecID;

	if( m_pFileOutFormat->video_codec != CODEC_ID_NONE && m_fCreateVideoStream ){
		m_pVideoSt = add_video_stream(vid);
		if(!m_pVideoSt) throw "Failed when create video stream!";
	}
	if( m_pFileOutFormat->audio_codec != CODEC_ID_NONE && m_fCreateAudioStream ){
		m_pAudioSt = add_audio_stream(aid);
		if(!m_pAudioSt)
			throw "Failed when create audio stream!";
	}

	if( av_set_parameters(m_pFormatContext, NULL) < 0 )
		throw "Failed when set format context parameters";

	if (!(m_pFileOutFormat->flags & AVFMT_NOFILE)) {
		if (url_fopen(&m_pFormatContext->pb,m_pMediaFile->m_strPath.GetString() , URL_WRONLY) < 0) {
			throw "Failed when open the output file!\n please check if there has enough disk space or \nthe file used by other process.";
		}
	}
	if (m_pVideoSt)
		if( !open_video() )
			throw "Failed when open video encoder! Please check your video steam parameters setting.";
	if(m_pAudioSt)
		if( !open_audio() )
			throw "Failde when open audio encoder! Please check your audio stream parameters setting.";
	av_write_header(m_pFormatContext);
	
	if(m_fCreateAudioStream && 
		(m_pProp->asampleRate != m_pSetting->m_nAudioSampleRate || 
		m_pProp->channels != m_pSetting->m_nAudioChannels))
	{
		m_pAudRmp = CreateLibAuResample();
		if(!m_pAudRmp) throw "Create audio sampler failed!";
		if(!m_pAudRmp->Init(m_pSetting->m_nAudioChannels, m_pSetting->m_nAudioSampleRate, m_pProp->channels, m_pProp->asampleRate))
			throw "Init Audio sampler failed!";
	}

	if( m_fCreateVideoStream && 
		(m_pProp->width != m_pSetting->m_nFrameWidth || 
		m_pSetting->m_nFrameHeight != m_pProp->height))
	{
		m_pImgRmp = CreateLibImgResample();
		if(!m_pImgRmp) throw "Alloc Resampler failed!";

		if(!m_pImgRmp->Init(m_pProp->width, m_pProp->height, m_pSetting->m_nFrameWidth, m_pSetting->m_nFrameHeight))
			throw "Init image sampler failed!";
	}	
		
	return true;
}



AVStream* CFFCore::add_audio_stream(int codec_id)
{
	if(!m_pFormatContext)return NULL;

	AVCodecContext *c;
	AVStream *st = NULL;
	
	st = av_new_stream(m_pFormatContext,1);
	if (!st) {
		throw "Failed when create a new audio stream.";		
	}

	c = st->codec;
	c->codec_id = (CodecID)codec_id;
	c->codec_type = CODEC_TYPE_AUDIO;


	c->bit_rate = m_pSetting->m_nAudioBitRate;
	c->sample_rate = m_pSetting->m_nAudioSampleRate;
	c->channels = m_pSetting->m_nAudioChannels;

	return st;
}

AVStream* CFFCore::add_video_stream(int codec_id)
{
	if(!m_pFormatContext)return NULL;

	AVCodecContext *c;
	AVStream *st = NULL;	
	st = av_new_stream(m_pFormatContext,0);
	
	if(!st) throw "Failed when create a new video stream!";

	c = st->codec;
	c->codec_id = (CodecID)codec_id;
	c->codec_type = CODEC_TYPE_VIDEO;

	
	c->bit_rate = m_pSetting->m_nVideoBitRate;	
	c->width = m_pSetting->m_nFrameWidth;
	c->height = m_pSetting->m_nFrameHeight;	
	c->time_base.den = m_pSetting->m_nFrameRate;
	c->time_base.num = m_pSetting->m_nFrameRateBase;
	c->gop_size = m_pSetting->m_nGopSize; 	
	c->pix_fmt = PIX_FMT_YUV420P;	
	
	if (c->codec_id == CODEC_ID_MPEG2VIDEO) {	
		c->max_b_frames = 2;
	}
	if (c->codec_id == CODEC_ID_MPEG1VIDEO){		
		c->mb_decision=2;
	}
	
	if(!strcmp(m_pFormatContext->oformat->name, "mp4") || !strcmp(m_pFormatContext->oformat->name, "mov") || !strcmp(m_pFormatContext->oformat->name, "3gp"))
		c->flags |= CODEC_FLAG_GLOBAL_HEADER;
	return st;
}



bool CFFCore::open_video()
{
	if(!m_pFormatContext || !m_pVideoSt) return false;

	AVCodec *codec;
	AVCodecContext *c =  m_pVideoSt->codec;	

	codec = avcodec_find_encoder(c->codec_id);
	if(!codec)throw "Can not find the video encoder!";

	
	int ret = 0;
	if ( (ret = avcodec_open(c, codec)) < 0) {
		throw "Failed when open video encoder!";
	}
	if(!m_pVidoeOutBuf){
		m_pVidoeOutBuf = (uint8_t*)av_malloc(MAX_VIDEO_OUT_SIZE);
		m_nVideoOutSize = MAX_VIDEO_OUT_SIZE;
		if(!m_pVidoeOutBuf) 
			throw "Failed when allocate memory!";
	}
	return true;
}





bool CFFCore::open_audio()
{
	if(!m_pFormatContext || !m_pAudioSt) return false;

	AVCodec *codec;
	AVCodecContext *c = m_pAudioSt->codec;	

	/* find the audio encoder */
	codec = avcodec_find_encoder(c->codec_id);
	if (!codec) {
		throw "Cannot find the audio encoder!";
	}

	/* open it */
	if (avcodec_open(c, codec) < 0) {
		throw "Failed when open audio encoder!";		
	}
	if(!m_pAudioOutBuf){
		m_pAudioOutBuf = (uint8_t*)av_malloc(MAX_AUDIO_OUT_SIZE);		
		if(!m_pAudioOutBuf)
			throw "Failed when allocate memory!";
	}
	if(!m_pAudioFifoBuf){
		m_pAudioFifoBuf = (uint8_t*)av_malloc(MAX_AUDIO_OUT_SIZE);
		if(!m_pAudioFifoBuf)
			throw "Failed When allocate memory!";
	}

	return true;
}


bool CFFCore::write_audio_frame(CAudioFrame* pFrame)
{	

	AVCodecContext *c;
	c = m_pAudioSt->codec;		
	short* pSample = NULL;
	int  datasize = 0;
	int frame_size = 0;

	if(!pFrame->GetEOS()){
		pSample = pFrame->GetDataPointer();
		datasize = pFrame->GetDataSize();
	}else {
		return true;
	}
	
	if(c->frame_size > 1){
		av_fifo_write(fifo, (uint8_t*)pSample, datasize);
		frame_size = c->frame_size * 2 * c->channels;
		while (av_fifo_read(fifo, m_pAudioFifoBuf, frame_size) == 0){
			int ret = avcodec_encode_audio(c, m_pAudioOutBuf, MAX_AUDIO_OUT_SIZE, (short*)m_pAudioFifoBuf);
			if(ret <= 0) return true;

			AVPacket pkt;
			av_init_packet(&pkt);
			pkt.size = ret;
			pkt.pts= av_rescale_q(c->coded_frame->pts, c->time_base, m_pAudioSt->time_base);
			pkt.flags |= PKT_FLAG_KEY;
			pkt.stream_index= m_pAudioSt->index;
			pkt.data= m_pAudioOutBuf;
			if (av_write_frame(m_pFormatContext, &pkt) != 0) {
				throw "There is an error when write audio frame!";	
			}
		}
	}else{
		switch(c->codec->id) {
		case CODEC_ID_PCM_S16LE:
		case CODEC_ID_PCM_S16BE:
		case CODEC_ID_PCM_U16LE:
		case CODEC_ID_PCM_U16BE:
			break;
		default:
			datasize = datasize >> 1;
			break;
		}

		int ret = avcodec_encode_audio(c, m_pAudioOutBuf, MAX_AUDIO_OUT_SIZE, (short*)m_pAudioFifoBuf);
		if(ret <= 0) return true;

		AVPacket pkt;
		av_init_packet(&pkt);
		pkt.size = ret;
		pkt.pts= av_rescale_q(c->coded_frame->pts, c->time_base, m_pAudioSt->time_base);
		pkt.flags |= PKT_FLAG_KEY;
		pkt.stream_index= m_pAudioSt->index;
		pkt.data= m_pAudioOutBuf;
		if (av_write_frame(m_pFormatContext, &pkt) != 0) {
			throw "There is an error when write audio frame!";	
		}
	}
	return true;
}
bool CFFCore::write_video_frame(CVideoFrame* pFrame)
{	
	int out_size, ret;
	AVCodecContext *c = m_pVideoSt->codec;	
	AVFrame* pSample = NULL;
	if(! pFrame->GetEOS()) 
		pSample = pFrame->GetAVFrame();
	pSample->pts = c->frame_number;

	if ( (m_pFormatContext->oformat->flags & AVFMT_RAWPICTURE) && pSample) {
		AVPacket pkt;
		av_init_packet(&pkt);

		pkt.flags |= PKT_FLAG_KEY;
		pkt.stream_index= m_pVideoSt->index;
		pkt.data= (uint8_t *)pSample;
		pkt.size= sizeof(AVPicture);
		ret = av_write_frame(m_pFormatContext, &pkt);
	} else {		
		out_size = avcodec_encode_video(c, m_pVidoeOutBuf, m_nVideoOutSize, pSample);
		if (out_size > 0) {
			AVPacket pkt;
			av_init_packet(&pkt);
			
			pkt.pts= av_rescale_q(c->coded_frame->pts, c->time_base, m_pVideoSt->time_base);		
			if(c->coded_frame->key_frame) 
				pkt.flags |= PKT_FLAG_KEY;
			pkt.stream_index= m_pVideoSt->index;
			pkt.data= m_pVidoeOutBuf;
			pkt.size= out_size;
			ret = av_write_frame(m_pFormatContext, &pkt);
		} else {
			ret = 0;
		}
	}
	if (ret != 0) {
		throw "Failed when write video frame!";		
	}
	return true;
}


void CFFCore::closeFormatContext()
{
	if(!m_pFormatContext) return;
	for (int i = 0; i < m_pFormatContext->nb_streams; i++)
	{		
		avcodec_close(m_pFormatContext->streams[i]->codec);
		av_free(m_pFormatContext->streams[i]->codec);
		av_free(m_pFormatContext->streams[i]);
	}	
}
	


bool CFFCore::WriteEnd()
{
	write_audio_end();
	write_video_end();

	if(!m_pFormatContext)
	av_write_trailer(m_pFormatContext);

	if (!(m_pFileOutFormat->flags & AVFMT_NOFILE)) {		
		url_fclose(&m_pFormatContext->pb);
	}
	return true;
}

bool CFFCore::write_video_end()
{
	int out_size, ret;
	if(!m_pVideoSt) return false;

	AVCodecContext *c = m_pVideoSt->codec;
	if ( m_pFormatContext->oformat->flags & AVFMT_RAWPICTURE )
	{
		return true;	
	} 
	else 
	{		
		out_size = avcodec_encode_video(c, m_pVidoeOutBuf, m_nVideoOutSize, NULL);		
		if (out_size > 0) {
			AVPacket pkt;
			av_init_packet(&pkt);
			if(c->coded_frame && c->coded_frame->pts != AV_NOPTS_VALUE)
			pkt.pts= av_rescale_q(c->coded_frame->pts, c->time_base, m_pVideoSt->time_base);
			if(c->coded_frame->key_frame) 
				pkt.flags |= PKT_FLAG_KEY;
			pkt.stream_index= m_pVideoSt->index;
			pkt.data= m_pVidoeOutBuf;
			pkt.size= out_size;
			ret = av_write_frame(m_pFormatContext, &pkt);
		} else {
			ret = 0;
		}
	}
	if (ret != 0) {
		throw("Write video end failed.");		
	}
	return true;

}

bool CFFCore::write_audio_end()
{
	if(!m_pAudioSt) return false;

	int out_size;
	out_size = avcodec_encode_audio(m_pAudioSt->codec, m_pAudioOutBuf, MAX_AUDIO_OUT_SIZE, NULL);
	if(out_size > 0)
	{
		AVPacket pkt;
		av_init_packet(&pkt);

		pkt.size= out_size;
		pkt.pts= av_rescale_q(m_pAudioSt->codec->coded_frame->pts, m_pAudioSt->codec->time_base, m_pAudioSt->time_base);
		pkt.flags |= PKT_FLAG_KEY;
		pkt.stream_index= m_pAudioSt->index;
		pkt.data= m_pAudioOutBuf;
		if (av_write_frame(m_pFormatContext, &pkt) != 0) {
			throw "There is an error when write audio frame!";	
		}
	}
	return true;
}


bool CFFCore::EncodeFrame(CMediaFrame* p)
{
	if(!p) throw "Null frame pointer! ";
	if(p->GetEOS())return true;

	if(p->GetType() == MEDIA_TYPE_AUDIO && m_fCreateAudioStream){

		CAudioFrame* pp = NULL;
		if(m_pAudRmp){
			m_pAudRmp->Process((CAudioFrame*)p, &pp, ((CAudioFrame*)p)->GetSampleCount());
		}else{
			pp = (CAudioFrame*)p;
		}
		write_audio_frame(pp);

	}else if(p->GetType() == MEDIA_TYPE_VIDEO && m_fCreateVideoStream){
		
		int nbframe = 0;
		LONGLONG rt0, rt1;
		p->GetTime(rt0, rt1);
		double x = double(rt0)/double(m_nFTP) - m_nTotalFrameNumber;

		if(x < -1.1) {
			nbframe = 0;
		}else if( x > 1.1 ){
			nbframe = int(x + 0.5);
		}	
		
		CVideoFrame* pp = NULL;
		if(m_pImgRmp){
			m_pImgRmp->Process((CVideoFrame*)p, &pp);
		}else{
			pp = (CVideoFrame*)p;
		}

		for (int i = 0; i < nbframe; i++)
		{
			write_video_frame(pp);
			m_nTotalFrameNumber++;
		}						
	}
	return true;
}

